import numpy as np
import os

dtype_emu = {
    np.float16: 1,
    np.float32: 2,
    np.int8: 3,
    np.int16: 4,
    np.int32: 5,
    np.complex64: 16,
    np.complex128: 17
}
def gen_iscomplexobj_data():
    # 创建输入和输出目录
    os.makedirs("./input", exist_ok=True)
    os.makedirs("./output", exist_ok=True)

    dtype = np.complex64
    # dtype = np.int32

    # 核间不均分，单核计算量非对齐
    # 定义目标形状：(17, 1023)
    target_shape = (8,2048)
    total_elements = target_shape[0] * target_shape[1]

    # 根据数据类型生成对应数据
    if np.issubdtype(dtype, np.complexfloating):
        # 生成复数类型
        real_part = np.random.uniform(-10, 10, total_elements).astype(np.float32)
        mask_nonzero_imag = np.random.random(total_elements) < 0.5
        imag_part = np.where(
            mask_nonzero_imag,
            np.random.uniform(-5, 5, total_elements).astype(np.float32),
            0.0
        )
        input_flat = real_part + imag_part * 1j
        input_flat = input_flat.astype(dtype)

    elif np.issubdtype(dtype, np.integer):
        # 生成整数类型
        # 修复: 检查是否为有符号整数的正确方法
        is_signed = np.issubdtype(dtype, np.signedinteger)

        if(is_signed):
            low, high = (-100, 100)
        else:
            low,high = (0, 200)
        
        input_flat = np.random.randint(low, high, total_elements, dtype=dtype)

    elif np.issubdtype(dtype, np.floating):
        # 浮点类型生成逻辑
        input_flat = np.random.uniform(-10, 10, total_elements).astype(dtype)

    else:
        raise ValueError(f"不支持的数据类型: {dtype}")

    # 重塑为目标形状
    input_x = input_flat.reshape(target_shape)

    is_complex_type = np.iscomplexobj(input_x)  # 标量布尔值
    golden = np.array([1 if is_complex_type else 0], dtype=np.int8)  # 形状(1,)

    # 生成tiling信息（算子切分配置）
    # 通常包含总元素数、数据类型标识等关键信息
    # 这里用：[总元素数, 数据类型标识(complex64=6)]（类型标识可自定义）
    tiling = np.array([
        total_elements,  # 第一个元素：总元素数（17*1023=17391）
        dtype_emu[dtype] # 第二个元素：数据类型标识（自定义complex64为5）
    ], dtype=np.uint32)  # 用uint32存储，适配多数算子的解析逻辑

    # 保存为二进制文件
    tiling.tofile("./input/input_tiling.bin")  # 新增tiling文件
    input_x.tofile("./input/input_x.bin")    # 类型：complex64，每个元素8字节（实部4+虚部4）
    golden.tofile("./output/golden.bin")     # 类型：bool，每个元素1字节（True=1，False=0）


if __name__ == "__main__":
    gen_iscomplexobj_data()